#include "pch.h"
#include "hooks.h"

NTSTATUS NtCreateFileHook(
	_Out_ PHANDLE FileHandle,
	_In_ ACCESS_MASK DesiredAccess,
	_In_ POBJECT_ATTRIBUTES ObjectAttributes,
	_Out_ PIO_STATUS_BLOCK IoStatusBlock,
	_In_opt_ PLARGE_INTEGER AllocationSize,
	_In_ ULONG FileAttributes,
	_In_ ULONG ShareAccess,
	_In_ ULONG CreateDisposition,
	_In_ ULONG CreateOptions,
	_In_reads_bytes_opt_(EaLength) PVOID EaBuffer,
	_In_ ULONG EaLength
)
{
	if (PsGetCurrentProcessId != 0)
	{
		//
		// We're going to filter for our "magic" file name.
		//
		if (ObjectAttributes &&
			ObjectAttributes->ObjectName &&
			ObjectAttributes->ObjectName->Buffer)
		{
			//
			// Unicode strings aren't guaranteed to be NULL terminated so
			// we allocate a copy that is.
			//
			PWCHAR ObjectName = (PWCHAR)ExAllocatePool(NonPagedPool, ObjectAttributes->ObjectName->Length + sizeof(wchar_t));
			if (ObjectName)
			{
				memset(ObjectName, 0, ObjectAttributes->ObjectName->Length + sizeof(wchar_t));
				memcpy(ObjectName, ObjectAttributes->ObjectName->Buffer, ObjectAttributes->ObjectName->Length);

				//
				// Does it contain our special file name?
				//
				if (wcsstr(ObjectName, L"testhook"))
				{
					//kprintf("[+] infinityhook: Denying access to file: %wZ.\n", ObjectAttributes->ObjectName);
					LogInfo("NtCreateFileHook :%wZ.\n", ObjectAttributes->ObjectName);
					ExFreePool(ObjectName);

					//
					// The demo denies access to said file.
					//
					return STATUS_ACCESS_DENIED;
					//return STATUS_NO_SUCH_FILE;
				}

				ExFreePool(ObjectName);
			}
		}
	}
	//
	// We're uninterested, call the original.
	//
	//
	return NtCreateFileOrig(FileHandle, DesiredAccess, ObjectAttributes, IoStatusBlock, AllocationSize, FileAttributes, ShareAccess, CreateDisposition, CreateOptions, EaBuffer, EaLength);
}

NTSTATUS
NtOpenProcessHook(
	_Out_ PHANDLE ProcessHandle,
	_In_ ACCESS_MASK DesiredAccess,
	_In_ POBJECT_ATTRIBUTES ObjectAttributes,
	_In_opt_ PCLIENT_ID ClientId
)
{
	if (PsGetCurrentProcessId != 0)
	{
		if (ClientId->UniqueProcess == 3284)
		{
			return STATUS_ACCESS_DENIED;
		}


	}
	return NtOpenProcessOrig(ProcessHandle, DesiredAccess, ObjectAttributes, ClientId);
}

NTSTATUS
ObReferenceObjectByHandleHook(
	_In_ HANDLE Handle,
	_In_ ACCESS_MASK DesiredAccess,
	_In_opt_ POBJECT_TYPE ObjectType,
	_In_ KPROCESSOR_MODE AccessMode,
	_Out_ PVOID *Object,
	_Out_opt_ POBJECT_HANDLE_INFORMATION HandleInformation
)
{
	KIRQL irql;
	NTSTATUS Result = STATUS_ACCESS_DENIED;
	//	irql = KeRaiseIrqlToDpcLevel();
	Result = ObReferenceObjectByHandleOrig(Handle, DesiredAccess, ObjectType, AccessMode, Object, HandleInformation);
	if (ObjectType == *PsProcessType)
	{
		if (PsGetProcessId(*Object) == 3284)
		{
			ObDereferenceObject(*Object);
			//	LogInfo("ObReferenceObjectByHandle Denied :)\n");
			Result = STATUS_ACCESS_DENIED;
		}

	}
	else
	{
		NTSTATUS Result = ObReferenceObjectByHandleOrig(Handle, DesiredAccess, ObjectType, AccessMode, Object, HandleInformation);
	}
	//	KeLowerIrql(irql);
	return Result;
}

NTSTATUS
NtAllocateVirtualMemoryHook(
	_In_ HANDLE ProcessHandle,
	_Inout_ _At_(*BaseAddress, _Readable_bytes_(*RegionSize) _Writable_bytes_(*RegionSize) _Post_readable_byte_size_(*RegionSize)) PVOID *BaseAddress,
	_In_ ULONG_PTR ZeroBits,
	_Inout_ PSIZE_T RegionSize,
	_In_ ULONG AllocationType,
	_In_ ULONG Protect
)
{

	return NtAllocateVirtualMemoryOrig(ProcessHandle, BaseAddress, ZeroBits, RegionSize, AllocationType, Protect);
}

NTSTATUS
NtQueryVirtualMemoryHook(
	_In_ HANDLE ProcessHandle,
	_In_opt_ PVOID BaseAddress,
	_In_ MEMORY_INFORMATION_CLASS MemoryInformationClass,
	_Out_writes_bytes_(MemoryInformationLength) PVOID MemoryInformation,
	_In_ SIZE_T MemoryInformationLength,
	_Out_opt_ PSIZE_T ReturnLength
)
{
	return NtQueryVirtualMemoryOrig(ProcessHandle, BaseAddress, MemoryInformationClass, MemoryInformation, MemoryInformationLength, ReturnLength);
}

NTSTATUS
NtUserFindWindowExHook(
	IN HWND hwndParent,
	IN HWND hwndChild,
	IN PUNICODE_STRING pstrClassName OPTIONAL,
	IN PUNICODE_STRING pstrWindowName OPTIONAL,
	IN DWORD dwType)
{
	LogInfo("ClassName: %wZ WindowName: %wZ", *pstrClassName, *pstrWindowName);
	return NtUserFindWindowExOrig(hwndParent, hwndChild, pstrClassName, pstrWindowName, dwType);
}
